import numpy as np
import pandas as pd
import matplotlib.pylab as plt
from kulfan_to_coord import CST_shape
from scipy.optimize import leastsq

df = pd.read_csv("dataset/ag09.csv").to_numpy()
x1 = df[:, 0]
y1 = df[:, 1]

global airfoil_CST
dz = 0
N = x1.shape[0]
print(N)

def objfunc(x, var_nums):
    var_nums = int(var_nums)
    wl = x[0:var_nums]
    wu = x[var_nums:2*var_nums]
    wl[0] = abs(wl[0])
    wu[0] = -wl[0]

    global airfoil_CST
    airfoil_CST = CST_shape(wl, wu, dz, N)
    coordinates = airfoil_CST.inv_airfoil_coor(x1)
    y2 = coordinates[:][1]

    return y2-y1

# =============================================================================
# Optimization
# =============================================================================

var_nums = 6
p_init = np.random.randn(2 * var_nums)  # 初始化
plsq = leastsq(func=objfunc, x0=p_init, args=var_nums)
wl_new = plsq[0][0:var_nums]
wu_new = plsq[0][var_nums:2*var_nums]

# error
dy = np.abs(objfunc(plsq[0], var_nums))
range1 = np.arange(32, 97, 1)
range2 = np.concatenate((np.arange(0, 32), np.arange(97, 129)))
print("x < 0.3:", dy[range1].any() < 0.0001)
print("x >= 0.3:", dy[range2].any() < 0.001)
print("x < 0.3 (max):", np.max(dy[range1]))
print("x >= 0.3 (max):", np.max(dy[range2]))

# plt.plot(range(len(dy)), dy)
# plt.show()

def plot():
    if N % 2 == 0:
        z = 0
    else:
        z = 1
    airfoil_CST2 = CST_shape(wl_new, wu_new, dz, N+z)
    coordinates = airfoil_CST2.airfoil_coor()
    x_coor = coordinates[0]
    y_coor = coordinates[1]
    ax = plt.subplot(111)
    ax.plot(x_coor, y_coor, 'g-', linewidth=4, label='CST')
    ax.plot(x1, y1, 'r', label='original')
    legend = ax.legend(loc='lower center', frameon=False)
    plt.xlabel('x/c')
    plt.ylabel('y/c')
    plt.ylim(ymin=-0.5, ymax=0.5)
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    ax.yaxis.set_ticks_position('left')
    ax.xaxis.set_ticks_position('bottom')
    plt.show()

# UNCOMMENT TO PLOT
# plot()

print('wl = ' + str(wl_new))
print('wu = ' + str(wu_new))


